from pycuda import gpuarray

from pycuda import driver as drv

import numpy as np
from pycuda.compiler import SourceModule
from simulation import NUMERIC_TYPE

from plugin_system import Plugin

import copy

#IMPLEMENT KERNEL HERE
kernel_src="""
__global__ void hw_integrator(int number_passed, float timestep, float* in_positions, float *in_velocities, float* masses, float* out_accelerations, float* potential_energy_array, float* out_positions, float* out_velocities){
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < number_passed) {
    
    //Do Accelerations Function
    float potential_energy=0.0f;
    float acc_x=0.0f;
    float acc_y=0.0f;
    float acc_z=0.0f;
    const float G = 2.959122082855910E-04f; 
    const float mass = masses[i];

    const float pos_x = in_positions[3*i];
    const float pos_y = in_positions[3*i+1];
    const float pos_z = in_positions[3*i+2];
    
    for (int j = 0; j < number_passed; j++){
        if(i != j) {
            const float rvx = pos_x - in_positions[3*j];
            const float rvy = pos_y - in_positions[3*j+1];
            const float rvz = pos_z - in_positions[3*j+2];
            const float x1 = rvx*rvx;
            const float y1 = rvy*rvy; 
            const float z1 = rvz*rvz;
            const float r2DOT = x1 + y1 + z1;
            const float r = sqrt(r2DOT);
            const float d = G * masses[j] / r;
            const float d3 = d / r2DOT;
            acc_x = acc_x - (rvx * d3);
            acc_y = acc_y - (rvy * d3); 
            acc_z = acc_z - (rvz * d3);
             
            potential_energy = potential_energy - ( d * mass );
        }
    }
    
    out_accelerations[3*i]=acc_x;
    out_accelerations[3*i+1]=acc_y;
    out_accelerations[3*i+2]=acc_z;
    potential_energy_array[i]=potential_energy;
    
    out_velocities[3*i] = in_velocities[3*i] + timestep * acc_x;
    out_velocities[3*i+1] = in_velocities[3*i+1] + timestep * acc_y;
    out_velocities[3*i+2] = in_velocities[3*i+2] + timestep * acc_z;
    out_positions[3*i] = in_positions[3*i] + timestep * out_velocities[3*i];
    out_positions[3*i+1] = in_positions[3*i+1] + timestep * out_velocities[3*i+1];
    out_positions[3*i+2] = in_positions[3*i+2] + timestep * out_velocities[3*i+2];
    }
}
"""
#rv is a vector
#r2 scalar
#
class HWVerlet(object):
    def __init__(self, timestep=None):
        self.timestep=timestep
        #gravitational constant
        self.G = NUMERIC_TYPE(2.959122082855910E-04)
        import pycuda.autoinit
        kernel_module = SourceModule(kernel_src)
        self.kernel_function = kernel_module.get_function("hw_integrator")

    def integrate(self, mysolar, step):
        positions = mysolar.positions
        velocities =mysolar.velocities
        masses = mysolar.masses
        #KICK
        #unpack positions and store into array that will be passed to GPU
        #name of
        names = mysolar.planet_names

        positions_array=np.zeros(len(positions)*3, np.float32)
        velocities_array=np.zeros(len(velocities)*3, np.float32)
        masses_array=np.zeros(len(masses), np.float32)

        for i, name in enumerate(names):
            positions_array[3*i: 3*i+3]=positions[name]
            velocities_array[3*i: 3*i+3]=velocities[name]
            masses_array[i]=masses[name]

        gpu_positions_array=gpuarray.to_gpu(positions_array)
        #unpack velocities and store into array that will be passed to GPU
        gpu_velocities_array=gpuarray.to_gpu(velocities_array)
        #unpack masses and store into array
        gpu_masses_array=gpuarray.to_gpu(masses_array)

        #accelerations  create
        accelerations_array=np.zeros(len(velocities) * 3, np.float32)
        gpu_accelerations_array=gpuarray.to_gpu(accelerations_array)
        #potential energy create
        potential_energy_array=np.zeros(len(velocities), np.float32)
        gpu_potential_energy_array=gpuarray.to_gpu(potential_energy_array)
        #output position array
        output_position_array=np.zeros(len(positions)*3, np.float32)
        gpu_output_position_array=gpuarray.to_gpu(output_position_array)
        #output velocities array
        output_velocity_array=np.zeros(len(positions)*3, np.float32)
        gpu_output_velocity_array=gpuarray.to_gpu(output_velocity_array)

        #set up kernel stuff
        threads_per_block = (128, 1, 1)
        number_passed=np.int32(len(names))
        gpu_timestep = np.float32(self.timestep)
        num_blocks = (len(names) / 128 + 1,1)
        #call kernel_function, one thread will execute EVERYTHING for ONE planet in system

        self.kernel_function(number_passed, gpu_timestep, gpu_positions_array, gpu_velocities_array, gpu_masses_array, gpu_accelerations_array,gpu_potential_energy_array, gpu_output_position_array, gpu_output_velocity_array, block=threads_per_block,grid=num_blocks) 

        #get stuf from the GPU output and return it
        accelerations_array=gpu_accelerations_array.get()
        potential_energy_array=gpu_potential_energy_array.get()
        output_velocity_array=gpu_output_velocity_array.get()
        output_position_array=gpu_output_position_array.get()

        pot_energy = 0.0
        out_positions = {}
        out_velocities = {}
        out_accelerations = {}
        for i, name in enumerate(names):
            out_accelerations[name]=accelerations_array[3*i:3*i+3]
            pot_energy += potential_energy_array[i]
            out_positions[name] = output_position_array[3*i:3*i+3]
            out_velocities[name] = output_velocity_array[3*i:3*i+3]

        kinetic_energy = self.kinetic_energy(velocities, masses)

        result = copy.deepcopy(mysolar)
        result.acceleration = out_accelerations
        result.positions = out_positions
        result.velocities = out_velocities
        result.potential_energy = pot_energy
        result.kinetic_energy = kinetic_energy
        result.simulated_time = step * self.timestep

        return result
  
    def kinetic_energy(self, velocities, masses):
        result=NUMERIC_TYPE(0.0)
        for planet, mass in masses.iteritems():
            v = velocities[planet]
            v2 = np.dot(v,v)
            result +=mass * v2
        return result*NUMERIC_TYPE(0.5)

class HWPluginAdapater(Plugin, HW):
    plugin_name = "HWIntegrator"
    parameter_definitions = {"timestep" : (NUMERIC_TYPE, True, "Timestep")}

    def __init__(self, *args, **user_params):
        params = self.configure_parameters(**user_params)
        HW.__init__(self, **params)

